import numpy as np

from matrixslow.core.node import Node


def fill_diagonal(to_be_filled, filter):
    '''
    将filter 矩阵填充在 to_be_filled 的对角线上
    :param to_be_filled:
    :param filter:
    :return:
    '''
    assert to_be_filled.shape[0] / filter.shape[0] == to_be_filled.shape[1] / filter.shape[1]

    n = int(to_be_filled.shape[0] / filter.shape[0])

    r, c = filter.shape
    for i in range(n):
        to_be_filled[i * r: (i + 1) * r, i * c: (i + 1) * c] = filter
    return to_be_filled


class Operator(Node):
    '''
    定义操作符抽象类
    什么都不做，只是起整理类继承结构的作用
    '''
    pass


class Add(Operator):
    '''
    （多个）矩阵加法
    '''

    def compute(self):
        # 构建一个和父节点中矩阵形状相同的累加器，并且初始为全零
        self.value = np.mat(np.zeros(self.parents[0].shape()))
        for parent in self.parents:
            self.value += parent.value

    def get_jacobi(self, parent):
        # 矩阵之和对其中任意一个矩阵的雅可比矩阵都是单位矩阵
        return np.mat(np.eye(self.dimension()))


class MatMul(Operator):
    '''
    矩阵乘法
    '''

    def compute(self):
        assert len(self.parents) == 2 and self.parents[0].shape()[1] == self.parents[1].shape()[0]

        self.value = self.parents[0].value * self.parents[1].value

    def get_jacobi(self, parent):
        '''
        将矩阵乘法视作映射，求映射对参与计算的矩阵的雅可比矩阵
        :param parent:
        :return:
        '''

        zeros = np.mat(np.zeros((self.dimension(), parent.dimension())))
        if parent is self.parents[0]:
            return fill_diagonal(zeros, self.parents[1].value.T)
        else:
            jacobi = fill_diagonal(zeros, self.parents[0].value)
            row_sort = np.arange(self.dimension()).reshape(self.shape()[::-1]).T.ravel()
            col_sort = np.arange(parent.dimension()).reshape(parent.shape()[::-1]).T.ravel()
            return jacobi[row_sort, :][:, col_sort]


class Step(Operator):
    '''
    施加阶跃函数操作的节点
    '''

    def compute(self):
        self.value = np.mat(np.where(self.parents[0].value >= 0.0, 1.0, 0.0))

    def get_jacobi(self, parent):
        np.mat(np.eye(self.dimension()))
        return np.zeros(np.where(self.parents[0].value >= 0.0, 0.0, -1.0))


class ScalarMultiply(Operator):
    '''
    用标量数（1*1矩阵）数乘一个矩阵
    '''
    def compute(self):
        assert self.parents[0].shape() == (1, 1)        # 第一个父节点必须是标量
        self.value = np.multiply(self.parents[0].value, self.parents[1].value)

    def get_jacobi(self, parent):
        assert parent in self.parents

        if parent is self.parents[0]:
            return self.parents[1].value.flatten().T
        else:
            return np.mat(np.eye(self.parents[1].dimension())) * self.parents[0].value[0, 0]

class Multiply(Operator):
    '''
    两个父节点的值是相同形状的矩阵，将它们对应位置的数相乘
    也叫做 数乘
    '''
    def compute(self):
        self.value = np.multiply(self.parents[0].value, self.parents[1].value)

    def get_jacobi(self, parent):
        if parent is self.parents[0]:
            return np.diag(self.parents[1].value.A1)
        else:
            return np.diag(self.parents[0].value.A1)


class Logistic(Operator):
    '''
    对向量的分量施加Logistic函数
    '''
    def compute(self):
        x = self.parents[0].value
        # 对父节点的每个分量施加Logistic
        self.value = np.mat(1.0 / (1.0 + np.power(np.e, np.where(-x > 1e2, 1e2, -x))))

    def get_jacobi(self, parent):
        return np.diag(np.mat(np.multiply(self.value, 1 - self.value)).A1)  # A1：展成一维数组

class SoftMax(Operator):
    '''
    SoftMax函数
    SoftMax([x1,x2,...,xk]) = (e^x1 / ∑e^xi, e^x2 / ∑e^xi, ...., e^xk / ∑e^xi)
    将输入向量的值都压缩到0，1之间，较大的分量得到接近于1的输出，较小的分量得到接近于0的输出，并且所有的输出之和是1
    可以把SoftMax的输出看作多个类别的概率分布
    '''

    @staticmethod
    def softmax(a):
        a[a > 1e2] = 1e2  # 防止指数过大
        ep = np.power(np.e, a)
        return ep / np.sum(ep)

    def compute(self):
        self.value = SoftMax.softmax(self.parents[0].value)

    def get_jacobi(self, parent):
        '''
        不实现SoftMax节点的get_jacobi函数
        训练时使用CrossEntropyWithSoftmax节点
        :param parent:
        :return:
        '''
        raise NotImplementedError("Don't use SoftMax's get_jacobi")


class ReLU(Operator):
    '''
    对矩阵的元素施加ReLU函数
    '''

    nslope = 0.1        # 负半轴的斜率

    def compute(self):
        self.value = np.mat(np.where(self.parents[0].value > 0.0, self.parents[0].value, self.nslope * self.parents[0].value))

    def get_jacobi(self, parent):
        return np.diag(np.where(self.parents[0].value.A1 > 0.0, 1.0, self.nslope))